using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;

using Atomic.Libraries.Mathematics;
using Atomic.Libraries.Plotting;
using Atomic.Thermodynamics.StateEquations;
using Atomic.Vasp;
using Atomic.Thermodynamics;
using Atomic.Structures;

namespace Atomic.Samples.Structures
{
	public static class Thermodynamics
	{
		public static void Sample1()
		{
			// Load volume and energy from VASP output files, one from each subdirectory.
			List<IVolumeEnergyPoint> points = new List<IVolumeEnergyPoint>();
			foreach (DirectoryInfo d in new DirectoryInfo("/home/mbjba/atomiclab/hcp").GetDirectories())
			{
				VaspResult r = new VaspResult(d);
				points.Add(new VolumeEnergyPoint(r.FinalStructure.Volume, r.Energy));
			}

			// Fit a BM4 equation of state (EOS) to the energies.
			IStateEquation bm = BirchMurnaghan4.Fit(points);

			// Plot points and fitted EOS using the method Energy defined by IStateEquation.
			new Gnuplot()
				.SetXLabel("Volume")
				.SetYLabel("Energy")
				.Plot(
					new DataPlot(points.Select(p => new PlaneVector(p.Volume, p.Energy)), LineColor.Black, PointType.Square),
					new FunctionPlot(v => bm.Energy(v), points.Min(p => p.Volume), points.Max(p => p.Volume), LineColor.Black)
				)
				.Pause()
				.Run();
		}

		public static void Sample2()
		{
			// Load again. This time using powerful LINQ expression. Store as VASP objects.
			VaspResult[] results = new DirectoryInfo("/home/mbjba/atomiclab/hcp")
				.GetDirectories()
				.Select(d => new VaspResult(d))
				.ToArray();

			IVolumeEnergyPoint[] points = results
				.Select(r => new VolumeEnergyPoint(r.FinalStructure.Volume, r.Energy))
				.ToArray();

			// Fit a BM4 equation of state (EOS) to the energies.
			BirchMurnaghan4 bm = BirchMurnaghan4.Fit(points);

			// Extract a structure object to calculate mass and number of atoms. Just take the first.
			Structure structure = results.First().FinalStructure;

			double scaling = 0.617;
			double gamma = (1.0 + bm.EquilibriumBulkModulusPressureDerivative) / 2.0 - 2.0 / 3.0;
			double mass = DebyeGruneisenPotential.StructureMassGeometric(structure);
			int atoms = structure.Sites.Count;

			// The everything is required to fit the Debye model is calculated.
			IHelmholtzPotential debye = new DebyeGruneisenPotential(bm.EquilibriumBulkModulus, bm.EquilibriumVolume, mass, atoms, scaling, gamma);

			// The total energy is elastic energy and vibrational free energy.
			IHelmholtzPotential pot = HelmholtzPotential.Add(HelmholtzPotential.Extend(bm), debye);

			// Plot with magnetic moment on the secondary axis.
			double t = 600.0;
			new Gnuplot()
				.SetXLabel("Volume")
				.SetYLabel("Energy")
				.SetYSecondaryLabel("Magnetic moment")
				.Plot(
					new DataPlot(points.Select(p => new PlaneVector(p.Volume, p.Energy)), LineColor.Black, PointType.Square),
					new DataPlot(results.Select(r => new PlaneVector(r.FinalStructure.Volume, r.MagneticMoment)), Axes.XYSecondary, LineColor.Blue, PointType.Circle),
					new FunctionPlot(v => bm.Energy(v), points.Min(p => p.Volume), points.Max(p => p.Volume), LineColor.Black, new PlotLabel("E_0")),
					new FunctionPlot(v => pot.FreeEnergy(t, v), points.Min(p => p.Volume), points.Max(p => p.Volume), LineColor.Red, new PlotLabel("F_ph"))
				)
				.Pause()
				.Run();
		}
	}
}
